# 使用AI自动生成结构化数据（一些虚拟数据等）,泛化AI模型
# pip install langchain_experimental

import os

from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_experimental.tabular_synthetic_data.openai import create_openai_data_generator
from langchain_experimental.tabular_synthetic_data.prompts import SYNTHETIC_FEW_SHOT_PREFIX, SYNTHETIC_FEW_SHOT_SUFFIX

from langchain_openai.chat_models.base import ChatOpenAI
# from langchain_core.pydantic_v1 import BaseModel, Field
from pydantic import BaseModel, Field


# 本地有clash代理，所以配置一下，不然一些依赖下载不下来（chroma）,或者将代理关闭后下载
os.environ['http_proxy'] = '127.0.0.1:7890'
os.environ['https_proxy'] = '127.0.0.1:7890'

os.environ["LANGCHAIN_TRACING_V3"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["LANGCHAIN_PROJECT"] = "langchain-community-demo"

os.environ["OPENAI_API_KEY"] = "sk-1dd16a258a73428d910d38c782e1c94f"

# deepseek-reasoner : DeepSeek-R1
# deepseek-chat : DeepSeek-V3
model_name = "deepseek-chat"
deepseek_api_key = "sk-1dd16a258a73428d910d38c782e1c94f"
yun_wu_api_key = "sk-y7KZj0I5wO10Zz3fsVPHj04QhXQtZjxl1FDm0SMEeqUBNqda"

## todo: deep seek目前不支持openAI的一些回调 function_call(可使用下面的模型，但是需要充值)
# model = ChatOpenAI(
#     model=model_name,
#     openai_api_key=deepseek_api_key,
#     openai_api_base='https://api.deepseek.com',
#     max_tokens=2048,
#     streaming=True,
#     temperature=0.7,
# )
model = ChatOpenAI(
    model='gpt-4-turbo',
    openai_api_key=yun_wu_api_key,
    openai_api_base='https://yunwu.ai/v1',
)


# 生成结构化数据 5个步骤


# 1、定义模型
class MedicalBilling(BaseModel):
    patient_name: str = Field(default=None, description="患者姓名")
    patient_id: str = Field(default=None, description="患者ID")
    patient_gender: str = Field(default=None, description="患者性别")
    patient_age: int = Field(default=None, description="患者年龄")
    patient_address: str = Field(default=None, description="患者地址")
    total_amount: float = Field(default=None, description="总金额")


# 2、提供一些样例给AI
examples = [
    {
        'example': "Patient Name: 张三, Patient Id: 12458845, Patient Gender: 男, Patient Age: 24, Patient Address: 济南, "
                   "Total Amount: 100"
    },
    {
        'example': "Patient Name: 李四, Patient Id: 5188551, Patient Gender: 女, Patient Age: 18, Patient Address: 北京, "
                   "Total Amount: 255.6"
    }
]

# 3、定义一个提示模版，用来指导AI生成结构化数据
openai_template = PromptTemplate(input_variables=['example'], template="请按照以下格式生成医疗账单信息：{example}")

prompt_template = FewShotPromptTemplate(
    prefix=SYNTHETIC_FEW_SHOT_PREFIX,
    examples=examples,
    suffix=SYNTHETIC_FEW_SHOT_SUFFIX,
    example_prompt=openai_template,
    input_variables=['subject', 'extra']
)

# 4、创建结构化数据生成器
generator = create_openai_data_generator(
    output_schema=MedicalBilling,
    prompt=prompt_template,
    llm=model
)

# 5、调用生成器
result = generator.generate(
    # 主题，和上面定义的key保持一致
    subject='医疗账单信息',
    # 额外信息
    extra='名字可以是随机的，最好使用比较生僻的名字',
    runs=5,
    # return_only_outputs=True
)

for record in result:
    print(record)


